package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"
	"io/ioutil"
	"log"
	"net/http"
	"runtime/debug"
	"strings"
	pd "zt/go-study/grpc/proto"
)

var _ pd.SearchServiceServer = (*SearchService)(nil)

type SearchService struct {
	pd.UnimplementedSearchServiceServer
}

func (s *SearchService) Search(ctx context.Context, r *pd.SearchRequest) (*pd.SearchResponse, error) {
	return &pd.SearchResponse{Response: r.GetRequest() + " Server"}, nil
}

const PORT = "9001"

func main() {
	certFile := "/Users/zhoutao/project/go-study/grpc/conf/server/server.pem"
	keyFile := "/Users/zhoutao/project/go-study/grpc/conf/server/server.key"
	caFile := "/Users/zhoutao/project/go-study/grpc/conf/ca.pem"
	c, err := GetTLSCredentialsByCA(certFile, keyFile, caFile)
	if err != nil {
	    log.Fatalf("GetTLSCredentialsByCA err: %v", err)
	}

	opts := []grpc.ServerOption{
		grpc.Creds(c),
		grpc_middleware.WithUnaryServerChain(LoggingInterceptor, RecoveryInterceptor),
	}

	server := grpc.NewServer(opts...)
	pd.RegisterSearchServiceServer(server, &SearchService{})

	mux := GetHTTPServeMux()

	// 为提供监听 HTTPS 服务的方法，重点的协议判断转发
	http.ListenAndServeTLS(":"+PORT,certFile, keyFile,http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			server.ServeHTTP(w, r)
		} else {
		    mux.ServeHTTP(w, r)
		}

		return
	}))
}

// GetHTTPServeMux 获取http服务
func GetHTTPServeMux() *http.ServeMux {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("eddycjy: go-grpc"))
	})

	return mux
}

func LoggingInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	log.Printf("gRPC method: %s, %v", info.FullMethod, req)
	resp, err := handler(ctx, req)
	log.Printf("gRPC method: %s, %v", info.FullMethod, resp)
	return resp, err
}

func RecoveryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
	defer func() {
	    if e := recover(); e != nil {
	        debug.PrintStack()
	        err = status.Errorf(codes.Internal, "Panic err: %v", e)
	    }
	}()

	return handler(ctx, req)
}

// GetTLSCredentialsByCA 通过ca获取tls证书
func GetTLSCredentialsByCA(certFile, keyFile, caFile string) (credentials.TransportCredentials, error) {
	// 从证书相关文件中读取和解析信息，得到证书公钥、密钥对
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		log.Fatalf("credentials.NewServerTLSFromFile err: %v", err)
		return nil, err
	}

	// 创建一个新的、空的 CertPool
	certPool := x509.NewCertPool()
	ca, err := ioutil.ReadFile(caFile)
	if err != nil {
		log.Fatalf("ioutil.ReadFile err: %v", err)
		return nil, err
	}

	// 尝试解析所传入的 PEM 编码的证书。如果解析成功会将其加到 CertPool 中，便于后面的使用
	if ok := certPool.AppendCertsFromPEM(ca); !ok {
		log.Fatalf("certPool.AppendCertsFromPEM err")
		return nil, err
	}

	// 构建基于 TLS 的 TransportCredentials 选项
	c := credentials.NewTLS(&tls.Config{
		Certificates: []tls.Certificate{cert}, // 设置证书链，允许包含一个或多个
		ClientAuth:   tls.RequireAndVerifyClientCert, // 要求必须校验客户端的证书
		ClientCAs:    certPool, // 设置根证书的集合，校验方式使用 ClientAuth 中设定的模式
	})

	return c, nil
}
